#include "model_base.h"
#include <thread>
#include <iostream>

ModelBase::ModelBase() : m_env(nullptr), m_session(nullptr)
{
}

ModelBase::~ModelBase()
{
}

bool ModelBase::LoadModel(const std::string &onnx_env_name, const std::string &onnx_path)
{
    try
    {
        // 宽字符
        std::wstring wstr = std::wstring(onnx_path.begin(), onnx_path.end());

        // 创建会话
        m_env = Ort::Env(ORT_LOGGING_LEVEL_ERROR, onnx_env_name.c_str());

        // 获取线程数
        int cpu_process_num = std::thread::hardware_concurrency();
        cpu_process_num /= 2;

        // 初始化设置
        Ort::SessionOptions session_options;
        // 设置线程数
        session_options.SetIntraOpNumThreads(cpu_process_num);
        // 优化级别
        session_options.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_BASIC);
        // 日志警告级别
        session_options.SetLogSeverityLevel(4);

        // 添加 CUDA 提供程序并检查状态
        Ort::ThrowOnError(OrtSessionOptionsAppendExecutionProvider_CUDA(session_options, 0));

        m_session = std::make_unique<Ort::Session>(m_env, wstr.c_str(), session_options); // 分配内存

        if (!m_session)
            return false;

        // PrintModelInfo(*m_session);

        return true;
    }
    catch (const Ort::Exception &e)
    {
        std::cerr << "ONNX Runtime error: " << e.what() << std::endl;
        return false;
    }
    catch (const std::exception &e)
    {
        std::cerr << "Standard exception: " << e.what() << std::endl;
        return false;
    }
    catch (...)
    {
        std::cerr << "Unknown exception occurred." << std::endl;
        return false;
    }
}

float ModelBase::LetterBox(const cv::Mat &input_mat, cv::Mat &output_mat,
                           std::vector<int> &padV,
                           const cv::Size &new_shape,
                           const cv::Scalar &color,
                           bool scale_up)
{
    cv::Mat convertMat;
    if (input_mat.channels() == 3)
    {

        cv::cvtColor(input_mat, convertMat, cv::COLOR_BGR2RGB);
    }

    else
    {
        cv::cvtColor(input_mat, convertMat, cv::COLOR_GRAY2RGB);
    }

    cv::Size shape = input_mat.size();
    // std::cout << "Original size: " << shape << std::endl; // 调试输出

    // 获取原图的宽高
    float w = shape.width;
    float h = shape.height;

    // 计算缩放比例
    float scalefactor = std::min(new_shape.height / h,
                                 new_shape.width / w);

    // 判断是否需要放大
    if (!scale_up)
    {
        scalefactor = std::min(scalefactor, 1.0f);
    }

    // 计数宽高
    int newUnpad[2]{(int)((float)w * scalefactor + 0.5), (int)((float)h * scalefactor + 0.5)};
    // std::cout << "New size: " << newUnpad[0] << "x" << newUnpad[1] << std::endl; // 调试输出

    cv::Mat tmp;

    if (w != newUnpad[0] || h != newUnpad[1])
    {
        cv::resize(convertMat, tmp, cv::Size(newUnpad[0], newUnpad[1]));
    }
    else
    {
        tmp = convertMat.clone();
    }

    // 初始化 output_mat
    // output_mat = cv::Mat::zeros(new_shape, input_mat.type());

    // 计算填充值
    int padW = (new_shape.width - newUnpad[0]) / 2;
    int padH = (new_shape.height - newUnpad[1]) / 2;

    padV.push_back(padW);
    padV.push_back(padH);

    // 填充值计算
    int top = int(std::round(padH - 0.1f));
    int bottom = int(std::round(padH + 0.1f));
    int left = int(std::round(padW - 0.1f));
    int right = int(std::round(padW + 0.1f));

    // std::cout << "Padding: top=" << top << ", bottom=" << bottom << ", left=" << left << ", right=" << right << std::endl; // 调试输出

    cv::copyMakeBorder(tmp, output_mat, top, bottom, left, right, cv::BORDER_CONSTANT, color);
    return 1.0 / scalefactor;
}
void ModelBase::PrintModelInfo(Ort::Session &session)
{
    // 获取模型输入输出节点
    size_t num_input_nodes = session.GetInputCount();
    size_t num_ouput_nodes = session.GetOutputCount();

    std::cout << "Number of input node is:" << num_input_nodes << std::endl;
    std::cout << "Number of output node is:" << num_ouput_nodes << std::endl;

    Ort::AllocatorWithDefaultOptions allocator;
    std::cout << std::endl;

    for (auto i = 0; i < num_input_nodes; i++)
    {
        std::cout << "The input op-name " << i << " is:" << session.GetInputNameAllocated(i, allocator) << std::endl;
    }

    for (auto i = 0; i < num_ouput_nodes; i++)
    {
        std::cout << "The output op-name " << i << " is:" << session.GetOutputNameAllocated(i, allocator) << std::endl;
    }

    // 输入输出维度
    for (auto i = 0; i < num_input_nodes; i++)
    {
        std::vector<int64_t> input_dims = session.GetInputTypeInfo(i).GetTensorTypeAndShapeInfo().GetShape();
        std::cout << std::endl
                  << "input" << i << "dims is:";

        for (auto j = 0; j < input_dims.size(); j++)
        {
            std::cout << input_dims[j] << " ";
        }
    }

    for (auto i = 0; i < num_ouput_nodes; i++)
    {
        std::vector<int64_t> output_dims = session.GetOutputTypeInfo(i).GetTensorTypeAndShapeInfo().GetShape();
        std::cout << std::endl
                  << "input" << i << "dims is:";
        for (auto j = 0; j < output_dims.size(); j++)
        {
            std::cout << output_dims[j] << " ";
        }
    }
}